# coding=utf-8
from __future__ import print_function, absolute_import
from gm.api import *
import datetime
import pandas as pd
import numpy as np
import pandas as pd
# 可以直接提取数据，掘金终端需要打开，接口取数是通过网络请求的方式，效率一般，行情数据可通过subscribe订阅方式
# 设置token， 查看已有token ID,在用户-密钥管理里获取
# set_token('bf739ec28ad507f391a3c3e99ed0365007da88f8')
set_token('ca9c4532786a122aa5ceb3fd726432862e046f91') # 135账户

# 获取目标日期date的A股代码（剔除停牌股、ST股、次新股（365天））
def get_normal_stocks(date,new_days=365):
    """
    获取目标日期date的A股代码（剔除停牌股、ST股、次新股（365天））
    :param date：目标日期字符串，格式为 "YYYY-MM-DD"
    :param new_days:新股上市天数，默认为365天
    """ 
    #=====================================================================
    date_obj = datetime.datetime.strptime(date, "%Y-%m-%d").replace(tzinfo=None).astimezone(tz=None)
    date = date_obj.replace(tzinfo=None).astimezone(tz=None)
    #=====================================================================
    
    # 先剔除退市股、次新股和B股
    # 调用 get_instrumentinfos 函数，获取所有股票的代码、上市日期和退市日期，
    # 并将这些信息以Pandas DataFrame的形式返回，存储在变量 df_code 中。
    # get_instrumentinfos - 查询交易标的基本信息
    df_code = get_instrumentinfos(sec_types=SEC_TYPE_STOCK, 
                                  fields='symbol, listed_date, delisted_date', 
                                  df=True)
    all_stocks = [
                    code        # 遍历前面筛选出的股票代码列表。
                    for code in df_code[
                        (df_code['listed_date'] <= date - datetime.timedelta(days=new_days)) & # 筛选出上市日期在 date 之前 new_days 天以上的股票。
                        (df_code['delisted_date'] > date + datetime.timedelta(days=31)) # 筛选出退市日期在 date 之后 31 天以上的股票。
                    ].symbol.to_list()  #将筛选后的 DataFrame 中的 symbol 列转换为列表。
                    if code[:6] != 'SHSE.9' and code[:6] != 'SZSE.2' #进一步筛选出股票代码前 6 个字符不是 SHSE.9 和 SZSE.2 的股票。
                    ]
    
    # 再剔除当前的停牌股和ST股
    history_ins = get_history_instruments(symbols=all_stocks, start_date=date, end_date=date, fields='symbol,sec_level, is_suspended', df=True)
    all_stocks = list(history_ins[(history_ins['sec_level']==1) & (history_ins['is_suspended']==0)]['symbol'])
    all_stocks = [code for code in all_stocks if code[:8]!='SHSE.688']
    # 将 all_stocks 列表中的所有元素用逗号 , 连接起来，形成一个字符串，并将这个字符串赋值给变量 all_stocks_str。
    all_stocks_str = ','.join(all_stocks)
    return all_stocks_str

# 读取通达信自选股函数
def process_TdxW_file(file_path):
    # 读取文本文件
    data = pd.read_csv(file_path, delimiter='\t', header=None, dtype=str)
    # 将DataFrame转换为字符串列表
    string_list = data.iloc[:, 0].tolist()
    # 替换第一个字符
    modified_string_list = [
        'SHSE.' + s[1:] if s.startswith(('1', '2')) else 'SZSE.' + s[1:] if s.startswith('0') else s
        for s in string_list
    ]
    # 将修改后的字符串列表拼接成一个字符串
    symbols_string = ','.join(modified_string_list)
    return symbols_string

# 自选股踢出次新股st等 函数
def find_intersection_stock(all_stock_str, symbol):
    # 将字符串拆分成列表
    a_list = all_stock_str.split(',')
    b_list = symbol.split(',')
    # 使用集合操作找到交集
    intersection = set(a_list) & set(b_list)
    # 将交集转换回字符串
    symbols_string = ','.join(intersection)
    return symbols_string

if __name__ == '__main__':
    file_path = 'E:\\自选股.EBK'  # 替换为你的文件路径
    date = "2024-07-19"
    symbol = process_TdxW_file(file_path) # 通达信自选股读取函数
    all_stock_str=get_normal_stocks(date,new_days=1000) #获取目标日期date的A股代码（剔除停牌股、ST股、次新股（365天）
    symbol=find_intersection_stock(all_stock_str, symbol) #自选股与筛选的A股交集（自选股中剔除停牌股、ST股、次新股（365天）
    print(symbol)
